package com.smallx.xpath.streaming;

import com.smallx.xpath.XPathMatcher;
import com.smallx.xpath.XPathMatcherFactory;
import com.smallx.xpath.parser.Expression;
import com.smallx.xpath.parser.LiteralExpression;
import com.smallx.xpath.parser.PairExpression;
import com.smallx.xpath.parser.NodeTestExpression;
import com.smallx.xpath.parser.NumberExpression;
import com.smallx.xpath.parser.SingletonExpression;
import com.smallx.xpath.parser.StepExpression;
import java.util.*;
import org.infoset.xml.Attribute;
import org.infoset.xml.Element;
import org.infoset.xml.Item;
import org.infoset.xml.Name;
import org.infoset.xml.XMLException;

/**
 * Compiles a subset of xpath to a event-oriented streaming expression.  The
 * subset supported is:
 *<pre>
 * xpath        := expr ( '|' expr )*
 * expr         := '/' step-expr | step-expr
 * step-expr    := element-step | text()
 * element-step := (qname | *) ( '/' element-step )?
 * qname        := name ( ':' name )?
 *</pre>
 * @author  R. Alexaner Milowski
 */
public class SubsetMatchCompiler {

   public static class AttributeValuePredicate {
      Name name;
      String value;
      boolean equals;
      
      AttributeValuePredicate(Name name,boolean equals,String value) {
         this.name = name;
         this.equals = equals;
         this.value = value;
      }
      
      AttributeValuePredicate(Name name) {
         this.name = name;
         this.equals = false;
         this.value = null;
      }
      
      public Name getName() {
         return name;
      }
      
      public String getValue() {
         return value;
      }
      
      public boolean isEquals() {
         return equals;
      }
   }
   
   /*
   static class HaltAndCatchFireExpression implements XPathExpression {
      public Iterator<Item> eval(Item context) 
         throws XMLException
      {
         throw new XMLException("Not implemented by the streaming compiler.");
      }
   }
   static class MatchEverythingExpression extends HaltAndCatchFireExpression {

      boolean doDocuments;
      MatchEverythingExpression(boolean doDocuments) {
         this.doDocuments = true;
      }
      public XPathMatcher newXPathMatcher() {
         return doDocuments ? (XPathMatcher)new MatchDocumentState() : (XPathMatcher)new MatchElementsState();
      }
      
   }
    */
   
   static class MatchDocumentState implements XPathMatcher {
      
      boolean matching;
      MatchDocumentState() {
         this.matching = false;
      }
      
      public Item match(Item current) 
         throws XMLException
      {
         if (!matching) {
            if (current.getType()==Item.ItemType.DocumentItem) {
               matching = true;
               //return MATCH;
               return current;
            }
         } else {
            if (current.getType()==Item.ItemType.DocumentEndItem) {
               matching = false;
            }
            //return MATCH;
            return current;
         }
         //return NO_MATCH;
         return null;
      }
      
      public void reset() {
         matching = false;
      }
      
   }

   static class MatchElementsState implements XPathMatcher {
      
      int level;
      MatchElementsState() {
         level = 0;
      }
      
      public Item match(Item current) 
         throws XMLException 
      {
         if (current.getType()==Item.ItemType.ElementItem) {
            level++;
         } else if (current.getType()==Item.ItemType.ElementEndItem) {
            level--;
            if (level==0) {
               return current;
               //return MATCH;
            }
         }
         return level!=0 ? current : null;
         //return level!=0 ? MATCH : NO_MATCH;
      }
      
      public void reset() {
         level = 0;
      }
      
   }
   
   interface State {
      State next(Item item);
      boolean isMatched();
   }
   
   static class ElementStepState implements State {
      Map transitions;
      State textMatch;
      ElementStepState wildcardMatch;
      List predicates;
      boolean matched;
      int positionRequired;
      
      class TextMatched implements State {
         public State next(Item item) {
            return doNext(item);
         }
         
         public boolean isMatched() {
            return true;
         }
         
      }
      
      ElementStepState() {
         transitions = new HashMap();
         textMatch = null;
         wildcardMatch = null;
         matched = false;
         predicates = null;
         positionRequired = -1;
      }
      
      public void addPredicate(AttributeValuePredicate predicate) {
         if (predicates==null) {
            predicates = new ArrayList();
         }
         predicates.add(predicate);
      }
      
      public int getPositionRequired() {
         return positionRequired;
      }
      
      public void setPositionRequried(int pos) {
         positionRequired = pos;
      }
      
      public State next(Item item) {
         return doNext(item);
      }
      State doNext(Item item) {
         switch (item.getType()) {
            case ElementItem:
            {
               Element e = (Element)item;
               Name ename = e.getName();
               //System.out.println(this+"doNext(): Checking element "+ename);
               State next = (State)transitions.get(ename);
               if (next==null) {
                  next = wildcardMatch;
               } else {
                  //System.out.println(this+"doNext(): Matched!");
                  if (predicates!=null) {
                     for (Iterator toCheck=predicates.iterator(); next!=null && toCheck.hasNext();) {
                        AttributeValuePredicate p = (AttributeValuePredicate)toCheck.next();
                        Attribute att = e.getAttributes().get(p.getName());
                        String valueToCheck = p.getValue();
                        if (att==null) {
                           next = null;
                        } else if (valueToCheck!=null && (p.isEquals() ? !att.getText().equals(valueToCheck) : att.getText().equals(valueToCheck))) {
                           next = null;
                        }
                     }
                  }
               }
               return next;
            }
            case CharactersItem:
               if (textMatch!=null) {
                  return textMatch;
               }
               break;
         }
         return null;
      }
      
      void addTransition(Name name,ElementStepState next) {
         transitions.put(name,next);
      }
      
      ElementStepState getTransition(Name name) {
         return (ElementStepState)transitions.get(name);
      }
      
      void matchText() {
         textMatch = new TextMatched();
      }
      
      ElementStepState matchWildcard() {
         if (wildcardMatch==null) {
            wildcardMatch = new ElementStepState();
         }
         return wildcardMatch;
      }
      
      public void setMatch() {
         matched = true;
      }
      
      public boolean isMatched() {
         return matched;
      }
      
   }
   
   /*
   static class SubsetEventExpression extends HaltAndCatchFireExpression  {

      State root;
      State other;
      
      SubsetEventExpression(State root,State other) {
         this.root = root;
         this.other = other;
      }
      
      public XPathMatcher newXPathMatcher() {
         return new SubsetXPathMatcher(root,other);
      }
      
   }*/
   
   static class SubsetXPathMatcher implements XPathMatcher {

      State root;
      State other;
      Stack currentStack;
      List countStore;
      int level;
      int index;
      boolean documentSeen;
      
      SubsetXPathMatcher(State root,State other) {
         this.root = root;
         this.other = other;
         this.level = 0;
         this.index = -1;
         this.currentStack = new Stack();
         this.countStore = new ArrayList();
         currentStack.push(Boolean.FALSE);
         this.documentSeen = false;
      }
      
      public Item match(Item item) 
         throws XMLException
      {
         if (level==0) {
            Item.ItemType type = item.getType();
            
            // If an element end, pop the current stack of matches
            if (type==Item.ItemType.ElementEndItem) {
               currentStack.pop();
            }
            
            // If it isn't an element or text, we can't match
            if (type!=Item.ItemType.ElementItem && type!=Item.ItemType.CharactersItem && type!=Item.ItemType.DocumentItem) {
               //return NO_MATCH;
               return null;
            }

            
            boolean isElement = type==Item.ItemType.ElementItem;

            int [] counts = null;
            Object top = currentStack.peek();
            List newCurrents = null;

            int last = 0;
            // If there are matches in play, check them first
            if (top!=null && top instanceof List) {
               
               List currents = (List)top;
               counts = index>=0 && index<countStore.size() ? (int [])countStore.get(index) : null;
               if (counts==null) {
                  counts = new int[currents.size()];
                  for (int i=0; i<counts.length; i++) {
                     counts[i] = 0;
                  }
                  countStore.add(counts);
               }
               
               last = currents.size();
               // We'll move down one in the count list as we have a match
               if (counts.length<last) {
                  int [] newcounts = new int[counts.length+currents.size()];
                  System.arraycopy(counts,0,newcounts,0,counts.length);
                  for (int i=counts.length; i<newcounts.length; i++) {
                     newcounts[i] = 0;
                  }
                  counts = newcounts;
                  countStore.set(index,counts);
               }
               // Iterate over the matches in play
               for (int i=0,len=currents.size(); i<len; i++) {
                  
                  // Check the next match
                  State current = (State)currents.get(i);
                  State nextCurrent = current.next(item);
                  
                  // Increment counters
                  if (nextCurrent!=null && isElement) {
                     counts[i]++;
                  }
                  
                  // Handle any positional requirement
                  if (nextCurrent!=null && nextCurrent instanceof ElementStepState) {
                     int posRequired = ((ElementStepState)nextCurrent).getPositionRequired();
                     if (posRequired>0 && posRequired!=counts[i]) {
                        nextCurrent = null;
                     }
                     //System.out.println("[prev matches] "+item+" index="+index+",posRequired="+posRequired+", counts["+i+"]="+counts[i]);
                  }
                     
                  // If there was a next state, modify the state
                  if (nextCurrent!=null) {
                     
                     // If it matched, we are done
                     if (nextCurrent.isMatched()) {
                        
                        // Only go one level for text, otherwise return the subtree
                        if (type==Item.ItemType.CharactersItem) {
                           return item;
                           //return MATCH;
                        } else {
                           level = 1;
                           //return MATCH;
                           return item;
                        }
                        
                     }
                     
                     // If we got here we should be an element and we have a new set of matches
                     if (isElement) {
                        if (newCurrents==null) {
                           newCurrents = new ArrayList();
                        }
                        newCurrents.add(nextCurrent);
                     }
                     
                  }
               }
            } 
            
            // Now that we have checked for old matches, check for new
            
            State current;
            // If we've seen the document, start with the other states
            if (documentSeen) {
               current = other.next(item);
               
            } else if (type==Item.ItemType.DocumentItem) { // If it is a document, check the root state
               current = root;
               isElement = true; // fake out the code below so that this state replaces the current one of false
               documentSeen = true;
            } else { // If the caller never passed a document, we assume they did
               current = other.next(item);
               documentSeen = true;
            }

            // Increment counters
            if (current!=null && isElement) {
               if (counts==null) {
                  index++;
                  counts = index>=0 && index<countStore.size() ? (int [])countStore.get(index) : null;
               }
               if (counts==null) {
                  //System.err.println("Allocating new counts at "+index);
                  counts = new int[1+last];
                  for (int i=0; i<counts.length; i++) {
                     counts[i] = 0;
                  }
                  if (index<countStore.size()) {
                     countStore.set(index,counts);
                  } else {
                     countStore.add(counts);
                  }
               } else if (counts.length<=last) {
                  int [] newcounts = new int[last+1];
                  System.arraycopy(counts,0,newcounts,0,counts.length);
                  for (int i=counts.length; i<newcounts.length; i++) {
                     newcounts[i] = 0;
                  }
                  counts = newcounts;
                  countStore.set(index,counts);
               }
               counts[last]++;
            }

            // Handle any positional requirement
            if (current!=null && current instanceof ElementStepState) {
               int posRequired = ((ElementStepState)current).getPositionRequired();
               //System.out.println("[new matches] "+item+" index="+index+",posRequired="+posRequired+", counts["+last+"]="+counts[last]);
               if (posRequired>0 && posRequired!=counts[last]) {
                  current = null;
                  index--;
               }
            }

            // If there was a state, we should add it to the matches
            if (current!=null) {
               
               // If we've matched, we are done
               if (current.isMatched()) {
                  
                  // Only go one level for text, otherwise return the subtree
                  if (type==Item.ItemType.CharactersItem) {
                     //return MATCH;
                     return item;
                  } else {
                     index--;
                     level = 1;
                     //return MATCH;
                     return item;
                  }
               }
        
               // If we got here we should be an element and we have a new set of matches
               if (isElement) {
                  if (newCurrents==null) {
                     newCurrents = new ArrayList();
                  }
                  newCurrents.add(current);
               }
            }
            
            // Since it is an element, we need to push a new stack context
            if (isElement) {
               currentStack.push(newCurrents==null ? (Object)Boolean.FALSE : (Object)newCurrents);
               if (newCurrents!=null) {
                  index++;
               }
            }
            //return currentStack.peek() instanceof Boolean ? NO_MATCH : MATCHING;
            return null;
         } else {
            switch (item.getType()) {
               case ElementItem:
                  level++;
                  break;
               case ElementEndItem:
                  level--;
                  break;
               case DocumentEndItem:
                  level--;
                  break;
            }
            return item;
            //return MATCH;
         }
      }
      
      public void reset() {
         currentStack.clear();
         currentStack.push(Boolean.FALSE);
         level = 0;
         countStore.clear();
         documentSeen = false;
      }
      
   }
   
   
   SubsetMatchCompiler() {
   }
   
   static XPathMatcherFactory compile(Expression ast) 
      throws XMLException 
   {
      // collapse AST and check for static errors via the following rules:
      // 1. Attribute targets are not allows and generate a static error
      // 2. A bare "/" expression collapses all other expressions
      // 3. Two paths that start the same collapse to the shorter path.
      // 4. "*" collapses all named axis steps.
      // 5. "text()" causes all text nodes to match.
      
      List rootSteps = new ArrayList();
      List otherSteps = new ArrayList();
      boolean matchAllTextNodes = false;
      boolean matchAllElements = false;
      boolean collapseAllSteps = false;

      Iterator<Expression> unionExprs = null;
      switch (ast.getOperator()) {
         case Union:
            unionExprs = ((List<Expression>)ast).iterator();
            break;
         case Root:
         case Path:
            unionExprs = Collections.singletonList(ast).iterator();
            break;
         default:
            throw new XMLException("Unsupported expression for compilation as matcher.");
      }
      ElementStepState rootState = null;
      ElementStepState otherState = new ElementStepState();
      while (unionExprs.hasNext() && !collapseAllSteps && !matchAllElements) {
         Expression top = unionExprs.next();
         if (top.getOperator()==Expression.Operator.Root) {
            if (collapseAllSteps || matchAllElements) {
               continue;
            }
            
            if (rootState==null) {
               rootState = new ElementStepState();
            }

            if (top instanceof SingletonExpression) {
               // Compile the FSM onto the root state
               compileFSM(rootState,(List<Expression>)((SingletonExpression)top).getExpression());
            } else {
               // A single "/" in the union
               collapseAllSteps = true;
               matchAllElements = false;
               matchAllTextNodes = false;
            }
            
         } else if (top.getOperator()==Expression.Operator.Path ) {
            // A non-rooted xpath (e.g. foo/bar)
            
            List<Expression> path = (List<Expression>)top;
            
            // If the size is 1, we could have * or text()
            if (path.size()==1 && path.get(0) instanceof StepExpression) {
               StepExpression step = (StepExpression)path.get(0);
               NodeTestExpression test = step.getTest();
               if (test.getType()==Item.ItemType.ElementItem) {
                  if (test.getName()==null) {
                     // It is a wildcard
                     matchAllElements = true;
                  } else {
                     // we just have a single named step for an element
                     compileFSM(otherState,path);
                  }
               } else if (test.getType()==Item.ItemType.CharactersItem) {
                  matchAllTextNodes = true;
               } else {
                  // we must have an unsupported test
                  throw new XMLException("Unsupported test type: "+test.getType());
               }
            } else {
               // we just compile the step path because it is multi step
               compileFSM(otherState,path);
            }
         } else {
            throw new XMLException("XPath expression not supported.");
         }
      }
      
      // NOTE: Subtree patterns that are masked by more general patterns
      // will be dropping the merge stage.  That is foo and foo/bar will
      // net only a 'foo' pattern in the merged FSM

      if (collapseAllSteps) {
         return new XPathMatcherFactory() {
            public XPathMatcher newMatcher() {
               return new MatchDocumentState();
            }
         };
      }
      
      if (matchAllElements) {
         return new XPathMatcherFactory() {
            public XPathMatcher newMatcher() {
               return new MatchElementsState();
            }
         };
      }
      
      final State matcherRootState = rootState;
      final State matcherOtherState = otherState;
      return new XPathMatcherFactory() {
         public XPathMatcher newMatcher() {
            return new SubsetXPathMatcher(matcherRootState,matcherOtherState);
         }
      };
   }
   
   static Name getAttributeName(List<Expression> predicatePath)
      throws XMLException
   {
      Expression attExpr = null;
      switch (predicatePath.size()) {
         case 1:
            attExpr = predicatePath.get(0);
            break;
         case 2:
            if (predicatePath.get(0).getOperator()!=Expression.Operator.Self) {
               throw new XMLException("Path expression "+predicatePath.get(0).getOperator()+" in predicate is not supported.");
            }
            attExpr = predicatePath.get(1);
            break;
         default:
            throw new XMLException("Path expression in predicate is not supported.");
      }
      if (attExpr.getOperator()!=Expression.Operator.Step) {
         throw new XMLException("Path expression in predicate is not supported.");
      }
      StepExpression attStep = (StepExpression)attExpr;
      if (attStep.getTest().getType()!=Item.ItemType.AttributeItem) {
         throw new XMLException("Path expression with node test "+attStep.getTest().getType()+" is not supported.");
      }
      return attStep.getTest().getName();
      
   }
   
   static void compileFSM(ElementStepState state,List<Expression> path) 
      throws XMLException
   {
      ElementStepState current = state;

      List<Expression> predicates = new ArrayList<Expression>();
      for (Iterator<Expression> exprs = path.iterator(); exprs.hasNext(); ) {
         Expression e = exprs.next();
         predicates.clear();
         while (e.getOperator()==Expression.Operator.Filter) {
            PairExpression f = (PairExpression)e;
            predicates.add(f.getSecond());
            e = f.getFirst();
         }
         Collections.reverse(predicates);
         if (e.getOperator()==Expression.Operator.Step) {
            StepExpression step = (StepExpression)e;
            NodeTestExpression test = step.getTest();
            if (test.getType()==Item.ItemType.CharactersItem) {
               // we can only have this at the end if the ast is built correctly
               current.matchText();
               continue;
            }
            if (test.getType()!=Item.ItemType.ElementItem) {
               throw new XMLException("Unsupport node test in path "+test.getType());
            }
            Name name = test.getName();
            ElementStepState predicated = current;
            if (name==null) {
               current = state.matchWildcard();
            } else {
               ElementStepState nextState = current.getTransition(name);
               if (nextState==null) {
                  nextState = new ElementStepState();
               }
               current.addTransition(name,nextState);
               current = nextState;
            }
            int position = -1;
            for (Expression p : predicates) {
               boolean equals = false;
               switch (p.getOperator()) {
                  // could be [@a]
                  case Path:
                  {
                     List<Expression> predicatePath =(List<Expression>)p;
                     Name attName = getAttributeName(predicatePath);
                     predicated.addPredicate(new AttributeValuePredicate(attName));
                  }
                     break;
                  // could be [@a='foo']
                  case Equals:
                     equals = true;
                  case NotEquals:
                  {
                     PairExpression comparison = (PairExpression)p;
                     Expression attExpr = comparison.getFirst();
                     Expression value = comparison.getSecond();
                     if (value.getOperator()==Expression.Operator.Path) {
                        attExpr = value;
                        value = comparison.getFirst();
                     }
                     
                     if (attExpr.getOperator()!=Expression.Operator.Path) {
                        throw new XMLException("Comparision is not with an attribute.");
                     }
                     if (value.getOperator()!=Expression.Operator.Literal) {
                        throw new XMLException("Attribute comparisons may only be with literals.");
                     }
                     Name attName = getAttributeName((List<Expression>)attExpr);
                     predicated.addPredicate(new AttributeValuePredicate(attName,equals,((LiteralExpression)value).getValue()));
                  }
                     
                     break;
                  // [1]
                  case Number:
                     if (position<1) {
                        position = (int)((NumberExpression)p).getValue();
                     } else {
                        int nextPos = (int)((NumberExpression)p).getValue();
                        if (nextPos!=1) {
                           throw new XMLException("Multiple position predicates will not match.  Second predicate is ["+nextPos+"]");
                        }
                     }
                     break;
                  default:
                     throw new XMLException("Predicate expression not supported: "+p.getOperator());
               }
            }
            if (position>=1) {
               current.setPositionRequried(position);
            }
         } else {
            throw new XMLException("Unsupport expression in path "+e.getOperator());
         }
      }
      current.setMatch();
   }
   
}
